import sys, getopt
import os
import struct
import ctypes
import binascii
import hashlib 
import zlib
from gmssl import sm3, func

# typedef enum
# {
#     OTA_HASH_TYPE_CHECKSUM = 0,
#     OTA_HASH_TYPE_XOR,
#     OTA_HASH_TYPE_CRC32,
#     OTA_HASH_TYPE_SHA1,
#     OTA_HASH_TYPE_SHA256,
#     OTA_HASH_TYPE_SM3,
#     OTA_HASH_TYPE_MAX
# } OTA_HASH_TYPE;
# typedef struct
# {
#     uint32_t magic;
#     uint32_t device;
#     uint32_t len;
#     uint32_t hash_type;
#     uint8_t hash_data[32];
# } user_fota_header_t;
# sizeof(user_fota_header_t)
HEAD_BACKUP_OFFSET = 48

def makeup_udp(type, firmwarefile, otafile):
    otafile_obj = open(otafile, 'wb+')

    firmware_obj = open(firmwarefile, 'rb')
    firmware_binary = firmware_obj.read()

    otafile_obj.seek(HEAD_BACKUP_OFFSET,0)
    otafile_obj.write(firmware_binary)

    fota_magic = 0xbeaf5aa5
    fota_device = 0x6750
    fota_len = os.path.getsize(firmwarefile)
    hash_type = type
    int_data = 0
    hash_data = ''
    print ("fota_len: ", fota_len)
    fota_head = ctypes.create_string_buffer(48)
    struct.pack_into('<I', fota_head, 0, fota_magic)
    struct.pack_into('<I', fota_head, 4, fota_device)
    struct.pack_into('<I', fota_head, 8, fota_len)
    struct.pack_into('<I', fota_head, 12, hash_type)

    if (type == 0):
        for val in list(firmware_binary):
            if (int_data > 4294967295):
                int_data = 0;
            int_data += val
        print ("int_data:", int_data)
        struct.pack_into('<I', fota_head, 16, int_data)
    elif (type == 1):
        for val in list(firmware_binary):
            int_data ^= val
        print ("int_data:", int_data)
        struct.pack_into('<I', fota_head, 16, int_data)
    elif (type == 2):
        # int_data = (binascii.crc32(firmware_binary) & 0xFFFFFFFF)
        int_data = zlib.crc32(firmware_binary) & 0xFFFFFFFF
        print ("int_data:", int_data)
        struct.pack_into('<I', fota_head, 16, int_data)
    elif (type == 3):
        sha1 = hashlib.sha1()
        sha1.update(firmware_binary)
        hash_data = sha1.hexdigest()
        print ("hash_data:", hash_data, "len:", len(hash_data))
        hash_arry = bytearray.fromhex(hash_data)
        for i in range(20):
            struct.pack_into('B', fota_head, 16+i, hash_arry[i])
    elif (type == 4):
        sha256 = hashlib.sha256()
        sha256.update(firmware_binary)
        hash_data = sha256.hexdigest()
        print ("hash_data:", hash_data, "len:", len(hash_data))
        hash_arry = bytearray.fromhex(hash_data)
        for i in range(32):
            struct.pack_into('B', fota_head, 16+i, hash_arry[i])
    elif (type == 5):
        hash_data = sm3.sm3_hash(func.bytes_to_list(firmware_binary))
        print ("hash_data:", hash_data, "len:", len(hash_data))
        hash_arry = bytearray.fromhex(hash_data)
        for i in range(32):
            struct.pack_into('B', fota_head, 16+i, hash_arry[i])
    else:
        print ("ERROR: no support the type:", type, "!")
        return


    otafile_obj.seek(0,0)
    otafile_obj.write(fota_head)

    otafile_obj.close()
    fota_len = os.path.getsize(otafile)

def help():
    print("---------help------------")
    print("一共4个参数")
    print("第一个参数：此脚本文件 pack_ota.py")
    print("第二个参数：签名type类型")
    print("           0-checksum 校验和校验")
    print("           1-xor 异或校验")
    print("           2-CRC32 校验")
    print("           3-SHA1 校验")
    print("           4-SHA256 校验")
    print("           5-SM3 校验")
    print("第三个参数：原始固件路径")
    print("第四个参数：目标OTA路径")
    print("")
    print("举例：")
    print("     python pack_ota.py 0 demo.bin update_sign.bin")
    print("     python pack_ota.py 1 demo.bin update_sign.bin")
    print("     python pack_ota.py 2 demo.bin update_sign.bin")
    print("     python pack_ota.py 3 demo.bin update_sign.bin")
    print("     python pack_ota.py 4 demo.bin update_sign.bin")
    print("     python pack_ota.py 5 demo.bin update_sign.bin")
    print("-------------------------")

def main(argv):
    argc = len(argv)
    if (argc < 4):
        help()
        return
    if (int(argv[1]) > 5):
        help()
        return
    print ("0:", argv[0], ",1:", argv[1], ",2:", argv[2], ",3:", argv[3])
    makeup_udp(int(argv[1]), argv[2], argv[3])
    print ("type:%d, pack %s -> %s Success!!!" %(int(argv[1]), argv[2],argv[3]))

if __name__ == '__main__':
    main(sys.argv)